import os
from datetime import datetime
import xarray as xr
import numpy as np

from config import conf
from create_nc import generate_new_nc_file
from log_own import logger

levels = [1000.0, 925.0, 850.0, 700.0, 500.0, 300.0, 250.0, 200.0, 50.0]


def produce_uv_wind(file_list):
    global init_time, time_units, lat_array, lon_array
    for lev in levels:
        u_list = []
        v_list = []
        t_list = []
        for file in file_list:
            try:
                if file.name.endswith("enfo-ef.grib2") or file.name.endswith("enfo-ep.grib2"):
                    return
                logger.info("开始生产 " + file.name + " " + str(lev) + "hPa的数据")
                u_ds = xr.open_dataset(file, filter_by_keys={'shortName': "u"})
                v_ds = xr.open_dataset(file, filter_by_keys={'shortName': "v"})
                u_da = u_ds["u"]
                v_da = v_ds["v"]
                try:
                    # 获取当前高度层的索引
                    level_array = u_ds["isobaricInhPa"].data
                    idx = np.where(level_array == lev)
                    # 取当前高度层的数据
                    u_data = u_da.data[idx][0]
                    v_data = v_da.data[idx][0]
                except IndexError:
                    logger.warning(str(lev) + "hPa高度数据不存在...")
                    continue
                # 获取 起报时间 和 预报时间
                init_time = file.stem.split('-')[0]  # 20230612000000
                time_units = datetime.strptime(init_time, "%Y%m%d%H%M%S").strftime("%Y-%m-%d %H:%M:%S")
                forecast_time = np.int16(file.stem.split('-')[1].replace("h", ""))  # 0
                # 获取经纬度数组
                lon_array = u_da["longitude"].data
                lat_array = u_da["latitude"].data

                u_list.append(u_data)
                v_list.append(v_data)
                t_list.append(forecast_time)
            except EOFError:
                logger.error(str(file) + "生产失败")

        if len(u_list) == 0:
            continue

        u_array = np.array(u_list)
        v_array = np.array(v_list)
        time_array = np.array(t_list)

        variables = {
            'longitude': {
                'array': lon_array,
                'nodata': -99999,
            },
            'latitude': {
                'array': lat_array,
                'nodata': -99999
            },
            'time': {
                'array': time_array,
                'units': "hours since " + time_units
            },
            'u': {
                'array': u_array,
                'nodata': -99999,
                'real_name': 'u'
            },
            'v': {
                'array': v_array,
                'nodata': -99999,
                'real_name': 'v'
            }
        }
        store_dir = os.path.join(conf["store_nc_dir"], "uv", init_time)
        if not os.path.exists(store_dir):
            os.makedirs(store_dir)
        nc_save_path = os.path.join(store_dir, ("uv" + "_" + str(lev) + ".nc"))
        if not os.path.exists(nc_save_path):
            # 生成nc文件
            generate_new_nc_file(variables, nc_save_path)


def produce_10uv_wind(file_list):
    global lon_array, lat_array, time_units, init_time
    u_list = []
    v_list = []
    t_list = []
    for file in file_list:
        try:
            if file.name.endswith("enfo-ef.grib2") or file.name.endswith("enfo-ep.grib2"):
                return
            logger.info("开始生产 " + file.name + " 的数据")
            u_ds = xr.open_dataset(file, filter_by_keys={'shortName': "10u"})
            v_ds = xr.open_dataset(file, filter_by_keys={'shortName': "10v"})
            u_da = u_ds["u10"]
            v_da = v_ds["v10"]

            # 取当前高度层的数据
            u_data = u_da.data
            v_data = v_da.data
            # 获取 起报时间 和 预报时间
            init_time = file.stem.split('-')[0]  # 20230612000000
            time_units = datetime.strptime(init_time, "%Y%m%d%H%M%S").strftime("%Y-%m-%d %H:%M:%S")
            forecast_time = np.int16(file.stem.split('-')[1].replace("h", ""))  # 0
            # 获取经纬度数组
            lon_array = u_da["longitude"].data
            lat_array = u_da["latitude"].data

            u_list.append(u_data)
            v_list.append(v_data)
            t_list.append(forecast_time)

        except EOFError:
            logger.error(str(file) + "生产失败")

    if len(u_list) == 0:
        return

    u_array = np.array(u_list)
    v_array = np.array(v_list)
    time_array = np.array(t_list)

    variables = {
        'longitude': {
            'array': lon_array,
            'nodata': -99999,
        },
        'latitude': {
            'array': lat_array,
            'nodata': -99999
        },
        'time': {
            'array': time_array,
            'units': "hours since " + time_units
        },
        '10u': {
            'array': u_array,
            'nodata': -99999,
            'real_name': '10u'
        },
        '10v': {
            'array': v_array,
            'nodata': -99999,
            'real_name': '10v'}
    }
    store_dir = os.path.join(conf["store_nc_dir"], "10uv", init_time)
    if not os.path.exists(store_dir):
        os.makedirs(store_dir)
    nc_save_path = os.path.join(store_dir, ("10uv" + ".nc"))
    if not os.path.exists(nc_save_path):
        # 生成nc文件
        generate_new_nc_file(variables, nc_save_path)